from django.conf import settings
from rest_framework import status
from rest_framework.viewsets import ViewSet
from rest_framework.response import Response
from rest_framework.decorators import action
from rest_framework.permissions import IsAuthenticated
from pymongo import MongoClient
from bson import ObjectId  # 用于处理 MongoDB 的 ObjectId


class TPAViewSet(ViewSet):
    permission_classes = [IsAuthenticated]

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.client = MongoClient(settings.MONGODB_URI)
        self.db = self.client[settings.MONGODB_DB]
        self.collection = self.db[settings.MONGODB_COLLECTION]

    @action(detail=False, methods=['post'])
    def save_data(self, request):
        try:
            data = request.data
            result = self.collection.insert_one(data)
            return Response({"status": "success", "message": "数据保存成功", "id": str(result.inserted_id)})
        except Exception as e:
            return Response({"status": "error", "message": str(e)})

    @action(detail=False, methods=['get'])
    def export_data(self, request):
        try:
            data = list(self.collection.find({}, {'_id': 0}))
            return Response(data)
        except Exception as e:
            return Response({"status": "error", "message": str(e)})

    @action(detail=False, methods=['get'])
    def view_data(self, request):
        try:
            problem_name = request.query_params.get('problem_name')

            # 获取分页参数
            page = int(request.query_params.get('page', 1))  # 默认第1页
            page_size = int(request.query_params.get('pageSize', 10))  # 默认每页10条
            # 构建查询条件
            query = {}
            if problem_name:
                query['problem'] = {'$regex': problem_name, '$options': 'i'}  # 模糊匹配，忽略大小写
            # 计算跳过的文档数量
            skip = (page - 1) * page_size

            # 查询数据总数
            total = self.collection.count_documents(query)

            # 分页查询数据
            data = list(self.collection.find(query).skip(skip).limit(page_size))

            # 将 ObjectId 转换为字符串
            formatted_data = []
            for item in data:
                item["_id"] = str(item["_id"])  # 将 _id 转换为字符串
                formatted_data.append(item)

            # 返回分页数据
            res_data = {
                'data': formatted_data,
                'total': total,  # 数据总数
                'page': page,  # 当前页码
                'pageSize': page_size,  # 每页条数
                'status': 'success',
                'message': '数据获取成功',
                'code': 200
            }
            return Response(res_data, status=status.HTTP_200_OK)
        except Exception as e:
            return Response({"status": "error", "message": str(e)})

    @action(detail=False, methods=['delete'])
    def delete_view(self, request):
        try:
            data_id = request.query_params.get('id')  # 从查询参数中获取 id
            if not data_id:
                return Response({"status": "error", "message": "ID 不能为空"})
            result = self.collection.delete_one({"_id": ObjectId(data_id)})
            if result.deleted_count == 0:
                return Response({"status": "error", "message": "未找到数据"})
            return Response({"status": "success", "message": "数据删除成功"})
        except Exception as e:
            return Response({"status": "error", "message": str(e)})

    @action(detail=False, methods=['get'])
    def get_data_by_id(self, request):
        try:
            data_id = request.query_params.get('id')  # 从查询参数中获取 id
            if not data_id:
                return Response({"status": "error", "message": "ID 不能为空"})

            # 根据 id 查询具体数据
            data = self.collection.find_one({"_id": ObjectId(data_id)})
            if not data:
                return Response({"status": "error", "message": "未找到数据"})

            # 将 ObjectId 转换为字符串
            data["_id"] = str(data["_id"])
            res_data = {
                'data': data,
                'status': 'success',
                'message': '数据获取成功',
                'code': 200
            }
            return Response(res_data, status=status.HTTP_200_OK)
        except Exception as e:
            return Response({"status": "error", "message": str(e)})

    @action(detail=False, methods=['post'])
    def update_data(self, request):
        try:
            data_id = request.data.get('id')  # 从请求体中获取 id
            if not data_id:
                return Response({"status": "error", "message": "ID 不能为空"})

            update_data = request.data.get('update', {})  # 获取更新数据
            if not update_data:
                return Response({"status": "error", "message": "更新数据不能为空"})

            # 移除 _id 字段，避免修改不可变字段
            if '_id' in update_data:
                del update_data['_id']

            # 根据 id 更新数据
            result = self.collection.update_one(
                {"_id": ObjectId(data_id)},
                {"$set": update_data}
            )

            if result.matched_count == 0:
                return Response({"status": "error", "message": "未找到数据"})

            return Response({"status": "success", "message": "数据更新成功"})
        except Exception as e:
            return Response({"status": "error", "message": str(e)})


